# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""export net together with checkpoint into air/mindir/onnx models"""
import argparse
import numpy as np
from mindspore import context, Tensor
from mindspore import load_checkpoint, load_param_into_net
from mindspore import export
from src.networks import define_G
from src.loss import InvertibleRescalingNetLoss

parser = argparse.ArgumentParser(description="IRN export")
parser.add_argument('--file_name', type=str, default='IRN_model', help='output file name prefix.')
parser.add_argument('--file_format', type=str, choices=['AIR', 'ONNX', 'MINDIR'], default='MINDIR',
                    help='file format')

parser.add_argument('--scale', type=int, default=2, choices=(2, 4),
                    help='Super resolution upscale factor.')
parser.add_argument('--resume_state', type=str,
                    default=None,
                    help='Path to the checkpoint.')
parser.add_argument("--device_id", type=int, default=0, help="device id, default: 0.")
parser.add_argument('--device_target', type=str, default='GPU', choices=("GPU", "CPU"),
                    help="Device target, support GPU, CPU.")

# Network specifications
parser.add_argument("--subnet_type", type=str, default='DBNet',
                    help="Type of subnet, default: DBNet.")
parser.add_argument("--in_nc", type=int, default=3,
                    help="Input channels, default: 3.")
parser.add_argument("--out_nc", type=int, default=3,
                    help="Output channels, default: 3.")
parser.add_argument("--block_num", type=int, default=[8, 8],
                    help="The block num, default: [8, 8].")
parser.add_argument("--G_scale", type=int, default=2,
                    help="Super resolution upscale factor, default: false.")
parser.add_argument("--G_init", type=str, default='xavier',
                    help="Init method of net, default: false.")

# Loss specifications
parser.add_argument("--pixel_criterion_forw", type=str, default='l2',
                    help="The pixel criterion of forward, default: l2.")
parser.add_argument("--pixel_criterion_back", type=str, default='l1',
                    help="The pixel criterion of backward, default: l1.")
parser.add_argument("--lambda_fit_forw", type=float, default=16,
                    help="lambda_fit_forw")
parser.add_argument("--lambda_rec_back", type=int, default=1,
                    help="lambda_rec_back")
parser.add_argument("--lambda_ce_forw", type=int, default=1,
                    help="lambda_ce_forw")
parser.add_argument("--gradient_clipping", type=int, default=10,
                    help="The values of gradient_clipping.")
parser.add_argument('--gaussian_scale', type=int, default=1,
                    help='The gaussian scale of net.')

if __name__ == '__main__':
    args = parser.parse_args()
    context.set_context(mode=context.GRAPH_MODE, device_id=args.device_id)

    # define network
    net = define_G(args)

    # define network with loss
    loss = InvertibleRescalingNetLoss(args, net)

    # loading resume state
    param_dict = load_checkpoint(args.resume_state)
    load_param_into_net(net, param_dict)

    loss.set_train(False)
    hr_image = Tensor(np.ones((1, 3, 144, 144), np.float32))
    lr_size = int(144 / args.scale)
    lr_image = Tensor(np.ones((1, 3, lr_size, lr_size), np.float32))
    inputs = [lr_image, hr_image]
    export(loss, *inputs, file_name=args.file_name, file_format=args.file_format)